# -*- coding: utf-8 -*-
"""
# @file name    : nn_layer_convolution.py
# @author       : QuZhang
# @date         : 2020-12-13 23:27
# @brief        : 学习卷积层
"""
import os
from tools.common_tools import transform_invert, set_seed
from PIL import Image
from torchvision.transforms import transforms
import torch.nn as nn
import matplotlib.pyplot as plt


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
path_tools = os.path.abspath(os.path.join(BASE_DIR, '..', '..', "tools", "common_tools.py"))
assert os.path.exists(path_tools), "{}不存在，请将common_tools.py文件放到 {}".format(path_tools, os.path.dirname(path_tools))
set_seed(3)

if __name__ == "__main__":
    # ================= load img ==============
    path_img = os.path.abspath(os.path.join(BASE_DIR, 'lena.png'))
    img = Image.open(path_img).convert('RGB')  # 0~255

    # from PIL.Image convert to Tensor
    img_transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    img_tensor = img_transform(img)
    img_tensor.unsqueeze_(dim=0)  # 添加batch维度，C*H*W to B*C*H*W

    # =============== create convolution layer ===============
    # =============== 正常卷积 Conv2d ==============
    # flag = True
    flag = False
    if flag:
        # (in_channels, out_channels, kernel_size, stride=1, padding=0, ...)
        conv_layer = nn.Conv2d(3, 1, 3)  # 创建卷积核
        nn.init.xavier_normal_(conv_layer.weight.data)  # 卷积核权重初始化

        img_conv = conv_layer(img_tensor)  # convolutional calculation

    # ============= 转置卷积 transposed ============
    # 转置卷积：用于上采样, 计算公式与正常卷积正好相反。
    flag = True
    # flag = False
    if flag:
        conv_layer = nn.ConvTranspose2d(3, 1, 3, stride=2)
        nn.init.xavier_normal_(conv_layer.weight.data)

        img_conv = conv_layer(img_tensor)

    # ================ visualization ===============
    print("卷积前img_tensor尺寸:{}\n卷积后img_conv尺寸:{}".format(img_tensor.shape, img_conv.shape))
    # print("*"*20)
    # print("img_conv:\n{}".format(img_conv))
    # print("-"*20)
    for i in range(img_conv.shape[0]):
        # print("#"*20)
        # print("img_conv[{}] size:{}\n{}".format(i, img_conv[i].shape, img_conv[i]))
        img_conv = transform_invert(img_conv[i], img_transform)
        img_raw = transform_invert(img_tensor.squeeze(), img_transform)
        plt.subplot(121).imshow(img_raw)
        plt.subplot(122).imshow(img_conv, cmap='gray')
        plt.show()

